import asyncio
from pyVim.connect import SmartConnectNoSSL, Disconnect
from pyVmomi import vim
import atexit
import os
from dotenv import load_dotenv

load_dotenv()

def WaitForTask(task):
    """
    同步等待vCenter任务完成，抛出详细异常信息
    :param task: vCenter任务对象
    :raises: 任务失败时抛出异常，包含详细错误信息
    """
    while task.info.state not in [vim.TaskInfo.State.success, vim.TaskInfo.State.error]:
        pass
    if task.info.state == vim.TaskInfo.State.error:
        error_msg = task.info.error.msg if hasattr(task.info.error, 'msg') else str(task.info.error)
        raise Exception(f"任务失败: {error_msg}")

def create_vm(host_ip: str, vm_name: str, cpu_count: int, memory_mb: int, disk_size_gb: int, datastore_name: str, network_name: str) -> dict:
    """
    MCP工具函数：创建虚拟机
    该函数连接vCenter服务器，基于传入的宿主机IP地址，虚拟机名称，CPU数量，内存大小（MB），硬盘大小（GB），硬盘存放的datastore名称，网络名称参数，
    在指定的ESXi宿主机上创建一个新的虚拟机。函数会自动查找对应的资源池、数据中心、存储和网络资源，配置虚拟机的硬件设备，包括CPU、内存、硬盘、网络适配器和默认CD/DVD驱动器。
    创建过程同步等待任务完成，返回一个字典，包含操作是否成功的标志success和详细的message信息。
    :param host_ip: ESXi宿主机IP地址，指定虚拟机创建的目标宿主机
    :param vm_name: 虚拟机名称，创建的虚拟机将使用该名称
    :param cpu_count: CPU数量，虚拟机分配的CPU核心数
    :param memory_mb: 内存大小（MB），虚拟机分配的内存容量，单位为兆字节
    :param disk_size_gb: 硬盘大小（GB），虚拟机硬盘容量，单位为千兆字节
    :param datastore_name: 硬盘存放的datastore名称，指定虚拟机硬盘存储位置
    :param network_name: 网络名称，虚拟机连接的网络名称
    :return: 创建结果字典，包含success（布尔值）和message（字符串）字段，指示操作结果和详细信息
    """
    result = {"success": False, "message": ""}
    try:
        vcenter_host = os.getenv("VCENTER_HOST")
        vcenter_user = os.getenv("VCENTER_USER")
        vcenter_password = os.getenv("VCENTER_PASSWORD")

        si = SmartConnectNoSSL(host=vcenter_host, user=vcenter_user, pwd=vcenter_password, port=443)
        atexit.register(Disconnect, si)

        content = si.RetrieveContent()

        # 查找指定宿主机
        host = content.searchIndex.FindByIp(None, host_ip, False)
        if not host:
            result["message"] = f"未能找到宿主机IP为 {host_ip}"
            return result

        # 获取资源池
        resource_pool = None
        if hasattr(host, "resourcePool"):
            resource_pool = host.resourcePool
        else:
            # 如果宿主机没有资源池，尝试获取集群资源池
            parent = host.parent
            if hasattr(parent, "resourcePool"):
                resource_pool = parent.resourcePool
        if not resource_pool:
            result["message"] = "未能找到资源池"
            return result

        # 获取datacenter
        datacenter = None
        parent = host.parent
        while parent:
            if isinstance(parent, vim.Datacenter):
                datacenter = parent
                break
            parent = getattr(parent, "parent", None)
        if not datacenter:
            result["message"] = "未能找到数据中心"
            return result

        # 获取datastore对象
        datastore = None
        for ds in datacenter.datastoreFolder.childEntity:
            if ds.name == datastore_name:
                datastore = ds
                break
        if not datastore:
            result["message"] = f"未能找到datastore {datastore_name}"
            return result

        # 获取网络对象
        network = None
        for net in datacenter.networkFolder.childEntity:
            if net.name == network_name:
                network = net
                break
        if not network:
            result["message"] = f"未能找到网络 {network_name}"
            return result

        # VM文件信息
        vmx_file = vim.vm.FileInfo(vmPathName=f'[{datastore_name}]')

        # 创建虚拟机配置
        config = vim.vm.ConfigSpec(
            name=vm_name,
            memoryMB=memory_mb,
            numCPUs=cpu_count,
            files=vmx_file,
            guestId='otherGuest',
            version='vmx-14'
        )

        # 添加SCSI控制器
        controller = vim.vm.device.VirtualDeviceSpec()
        controller.operation = vim.vm.device.VirtualDeviceSpec.Operation.add
        controller.device = vim.vm.device.ParaVirtualSCSIController()
        controller.device.busNumber = 0
        controller.device.sharedBus = vim.vm.device.VirtualSCSIController.Sharing.noSharing
        controller.device.key = 1000

        # 添加硬盘
        disk = vim.vm.device.VirtualDeviceSpec()
        disk.operation = vim.vm.device.VirtualDeviceSpec.Operation.add
        disk.fileOperation = vim.vm.device.VirtualDeviceSpec.FileOperation.create
        disk.device = vim.vm.device.VirtualDisk()
        disk.device.capacityInKB = disk_size_gb * 1024 * 1024
        disk.device.key = 0
        disk.device.unitNumber = 0
        disk.device.controllerKey = 1000
        disk.device.backing = vim.vm.device.VirtualDisk.FlatVer2BackingInfo()
        disk.device.backing.diskMode = 'persistent'
        disk.device.backing.thinProvisioned = True
        disk.device.backing.datastore = datastore

        # 添加网络适配器
        nic = vim.vm.device.VirtualDeviceSpec()
        nic.operation = vim.vm.device.VirtualDeviceSpec.Operation.add
        nic.device = vim.vm.device.VirtualVmxnet3()
        nic.device.deviceInfo = vim.Description()
        nic.device.backing = vim.vm.device.VirtualEthernetCard.NetworkBackingInfo()
        nic.device.backing.network = network
        nic.device.backing.deviceName = network_name
        nic.device.connectable = vim.vm.device.VirtualDevice.ConnectInfo()
        nic.device.connectable.startConnected = True
        nic.device.connectable.allowGuestControl = True

        config.deviceChange = [controller, disk, nic]

        # 添加默认CD/DVD驱动器
        cdrom = vim.vm.device.VirtualDeviceSpec()
        cdrom.operation = vim.vm.device.VirtualDeviceSpec.Operation.add
        cdrom.device = vim.vm.device.VirtualCdrom()
        cdrom.device.key = -1
        # 添加IDE控制器，如果没有则创建一个
        ide_controller = None
        for dev in config.deviceChange:
            if isinstance(dev.device, vim.vm.device.VirtualIDEController):
                ide_controller = dev.device
                break
        if not ide_controller:
            ide_spec = vim.vm.device.VirtualDeviceSpec()
            ide_spec.operation = vim.vm.device.VirtualDeviceSpec.Operation.add
            ide_spec.device = vim.vm.device.VirtualIDEController()
            ide_spec.device.key = 200
            ide_spec.device.busNumber = 0
            ide_spec.device.controllerKey = 100
            config.deviceChange.append(ide_spec)
            ide_controller = ide_spec.device
        cdrom.device.controllerKey = ide_controller.key
        cdrom.device.unitNumber = 0
        cdrom.device.backing = vim.vm.device.VirtualCdrom.IsoBackingInfo()
        cdrom.device.connectable = vim.vm.device.VirtualDevice.ConnectInfo()
        cdrom.device.connectable.startConnected = True
        cdrom.device.connectable.allowGuestControl = True
        config.deviceChange.append(cdrom)

        # 创建虚拟机任务，指定host确保虚拟机创建在指定宿主机
        vm_folder = datacenter.vmFolder
        # Remove unused placement spec, pass host and pool directly
        task = vm_folder.CreateVM_Task(config=config, pool=resource_pool, host=host)
        try:
            WaitForTask(task)
            result["success"] = True
            result["message"] = f"虚拟机 {vm_name} 创建成功"
        except Exception as e:
            result["message"] = f"创建虚拟机失败: {str(e)}"

    except Exception as e:
        result["message"] = f"创建虚拟机失败: {str(e)}"

    return result

